import atexit

from .impl.MySQLDriver import MySQLDB
from .impl.PostGreSQLDriver import PostGreSQLDB
from .impl.DBDriver import DB
from config import yaml_config, DatabaseSetting


def get_db_driver() -> DB:
    db: DB = None
    database_config: DatabaseSetting = yaml_config.database_setting
    db_type: str = database_config.type
    if db_type == 'mysql':
        db = MySQLDB()
    elif db_type == 'postgresql':
        db = PostGreSQLDB()
    else:
        raise Exception('数据库类别信息未设置(database.type)')
    return db


db: DB = get_db_driver()



@atexit.register
def close_db():
    db.close()
    print("数据库链接已关闭")